import torch
import torchvision.transforms as transforms
import sys
sys.path.append("models")
sys.path.append("utils")
import matplotlib.pyplot as plt
import numpy as np
import  tqdm 
from PIL import Image
import os
# from class_model import VGG19
from ghostnetv2_torch import ghostnetv2
from ShuffleNetv2_torch import ShuffleNetV2
import argparse
import cv2

def make_parser():
    parser = argparse.ArgumentParser(description="training config", add_help= False)

    parser.add_argument("--image_root", default='/home/cheng/data/workspace/part-time/ear/ff', type=str)
    parser.add_argument("--weight", default='output/ear_class2/2024_01_14_10_21_33/best.pth', type=str)
    parser.add_argument("--num_classes", default=3, type=int)
    parser.add_argument("--max_epoch", default=200, type=int)
    parser.add_argument("--batch_size", default=64, type=int)
    parser.add_argument("--lr_min", default=0.0001, type=float)
    parser.add_argument("--lr_max", default=0.01, type=float)
    parser.add_argument("--input_size", default=[576, 576], type=int, nargs='+')

    parser.add_argument("--load_train", default='', type=str)
    parser.add_argument("--resume_train", default='', type=str)
    parser.add_argument("--mean", default=[0.47627547, 0.3684924, 0.32831427], type=float, nargs='+')
    parser.add_argument("--std", default=[0.32356596, 0.26539913, 0.2421861], type=float, nargs='+')
    parser.add_argument("--num_workers", default=4, type=int)
    parser.add_argument("--base_lr", default=0.01, type=float)

    parser.add_argument("--output_name", default='', type=str)
    return parser

def main(args):
    from preproc import static_resize
    transform = transforms.Compose([
                            static_resize(args.input_size),
                            # transforms.RandomAffine(degrees=15,scale=(0.8,1.5)),
                            transforms.ToTensor(),
                            transforms.Normalize(args.mean, args.std)
                            ])

    # 加载模型
    model = ghostnetv2(num_classes=args.num_classes, 
                        width=1, 
                        dropout=0.2)
    # model = ShuffleNetV2(num_classes=2)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    # if torch.cuda.device_count() > 1:
    #     model = nn.DataParallel(model)

    model.to(device)
    # model = model.cuda()
    model.load_state_dict(torch.load(args.weight, map_location='cpu')['model'])       

    # 测试
    # id_list = []
    # pred_list = []
    test_files = os.listdir(args.image_root)
    model.eval()
    map_dict = {'call': 0, 'chat': 1, 'photo': 2}
    label_dict = {0: 'MTC', 1: 'CSOM', 2: 'normal'}
    with torch.no_grad():
        for file in tqdm.tqdm(test_files):
            img = Image.open(os.path.join(args.image_root, file))
            img = transform(img)
            # img.save("re.png")
            img = img.unsqueeze(0)
            img = img.to(device)
            # print(img)
            outs = model(img)
            print(outs.cpu().numpy())
            _, prediction = torch.max(outs, 1)
            # label = map_dict[file.split('_')[0]]
            # print("prediction = ", int(prediction.cpu().numpy()))
            # if int(prediction.cpu().numpy()) != label:
            cv_image = cv2.imread(os.path.join(args.image_root, file))
            # txt = label_dict
            cv2.putText(cv_image, str(label_dict[prediction.item()]), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            cv2.imwrite(os.path.join("/home/cheng/data/workspace/part-time/ear/res/", file), cv_image)
            # _predict = np.array(prediction)
            # _predict = np.where(_predict>0.5, 1, 0)
            # print(_id, _predict[0])
            # id_list.append(_id)
            # pred_list.append(_predict)
            # print(file, prediction.item())
    # res = pd.DataFrame({
    #     'id':id_list,
    #     'label':pred_list
    # })

    # res.sort_values(by='id', inplace=True)
    # res.reset_index(drop=True, inplace=True)
    # res.to_csv('submission.csv', index=False)

    # res.head(10)

    # import random

    # class_dict = {0:'cat', 1:'dog'}
    # fig, axes = plt.subplots(2, 5, figsize=(20,12), facecolor='w')

    # for ax in axes.ravel():
    #     i = random.choice(res['id'].values)
    #     label = res.loc[res['id']==i, 'label'].values[0]
    #     img = Image.open('../data/test/'+str(i)+'.jpg')
    #     ax.set_title(class_dict[label[0]])
    #     ax.imshow(img)

if __name__ == "__main__":
    args = make_parser().parse_args()
    main(args)